/**
  CNOK project, Anyang Normal University, IMP-CAS
  \class TAMOMDIS
  \brief This is a global class, to take control of the whole flow of the
  program. It is also responsible for generating various momentum distributions.
  \author SUN Yazhou, asia.rabbit@163.com
  \since 2020/07/09
  \date Last modified: 2020/10/05 by SUN Yazhou
  \copyright 2020-2023 SUN Yazhou
  \copyright CNOK project, Anyang Normal University, IMP-CAS
*/

#include <iostream>
#include <iomanip>
#include <cstring>
#include <sstream>
#include "config.h"
#ifdef CATCH_CNOK
#include "catch.h"
#endif
#include "TAMOMDIS.h"
#include "TAMOMDIS_M.h"
#include "TASMatrix.h"
#include "TABound.h"
#include "TAException.h"
#include "TAParaManager.h"
#include "TAFNN.h"
#include "TAYaml.h"
#include "TAConvolve.h"
#include "TAFileIO.h"
#include "TANucleus.h"
#include "TANuMom.h"
#include "TAMath.h"
#include "TAAME.h"

#define copy TAFileIO::Copy

using std::setw;
using std::setprecision;
using std::fixed;

static const TAAME *AME = TAAME::Instance(); // atomic mass evalution - 2021

TAMOMDIS::TAMOMDIS(int argc, char *argv[]) : fBound(nullptr), fSc(nullptr),
fSv(nullptr), fMOM_M(nullptr), fConv(nullptr){
  fIsProtonRemoval = -1; // -1: undefined; 1: true (zv: 1); 0: false (zv: 0)
  fNmom = TAMOMDIS_M::kNmom;
  fHasConfigured = false;
  fExCS = false;
  // set the basedir first
  const TAYaml bd("config/basedir.yaml");
  PARA->SetBaseDir(bd["basedir"].as<string>());
  fExpResCfgFile = "config/expres.yaml";

  GetOpt(argc, argv); // parse the command options
  PARA->Read(fCfgFile);
  if(!fIsBatch) Configure();
} // end ctor

// ctor for single mode
// note that if choosing this ctor, one has to call PARA->SetBaseDir(..)
// outside this ctor and before the calling of TAMOMDIS::Configure()
TAMOMDIS::TAMOMDIS(const string &cfgfile, bool isconv, bool isbound, bool ismom,
bool exCS) : fBound(nullptr), fSc(nullptr), fSv(nullptr), fMOM_M(nullptr), fConv(nullptr){
  fIsProtonRemoval = -1; // -1: undefined; 1: true (zv: 1); 0: false (zv: 0)
  fNmom = TAMOMDIS_M::kNmom;
  fHasConfigured = false; fIsBatch = false;
  fIsBound = isbound; fIsMom = ismom; fIsConv = isconv;
  fExCS = fIsBound ? false : exCS;
  fCfgFile = cfgfile;
  fExpResCfgFile = "config/expres.yaml";
  PARA->Read(fCfgFile);
  if(!fIsBatch) Configure();
} // end ctor

TAMOMDIS::~TAMOMDIS(){
  if(fSc) delete fSc;
  if(fSv) delete fSv;
  if(fBound) delete fBound;
  if(fMOM_M) delete fMOM_M;
  if(fConv) delete fConv;
  for(auto *t : fnv) if(t) delete t;
} // end dctor

// parse command options
void TAMOMDIS::GetOpt(int argc, char *argv[]){
  fIsBound = false; fIsMom = false; fIsBatch = false; fIsConv = false;
  fCfgFile = "";
  // opterr = 0;
  int ch; string batchYaml;
  if(1 == argc) PromptHelp();
	while((ch = getopt(argc, argv, ":m::b:r")) != -1){
		switch(ch){
      // calculate momdis
			case 'm': fIsMom = true; if(optarg && 'c' == optarg[0]) fIsConv = true; break;
			case 'b': fIsBatch = true; batchYaml = optarg; break; // batch mode
      case 'r': fIsBound = true; break; // only calculate Rl(r) and exit
			default: PromptHelp();
		} // end switch
	} // end while
  if(fIsBound && fIsBatch) TAExcp::Error("TAMOMDIS",
      "GetOpt: -b (testing Rl) and -r (batch mode) cannot be used simultaneously.");

  // enter single configuration mode
  if(argv[optind]){
    if(TAMath::endwith(argv[optind], ".yaml"))
      TAExcp::Error("TAMOMDIS", "GetOpt: single file with .yaml, input error.");
    char config[128]; sprintf(config, "%s.yaml", argv[optind]);
    fCfgFile = config;
  } // end if
  if(fCfgFile.empty()){
    if(batchYaml.empty())
      TAExcp::Error("TAMOMDIS", "GetOpt: empty arguments.");
    fCfgFile = batchYaml; // batch mode
  } // batch mode and single mode cannot co-exist
  else // single mode
    if(!batchYaml.empty()) TAExcp::Error("TAMOMDIS", "GetOpt: abnormal input.");
} // end member function GetOpt

// show manual for input rules
void TAMOMDIS::PromptHelp(){
	cout << "\n          -----------------\033[32;1m CNOK\033[0m";
  cout << " USER MANUL ----------------\n";
	cout << "Usage:\n";
	cout << "\t./mom [yamlFile, name without '.yaml' for single configuration]\n";
  cout << "\n\t\t[-m[c]] [-b <yamlFile>] [-r]" << endl;
  cout << "\n\n\t\t-m: calculate momentum distributions, (arg) c(optional, `-mc`):";
  cout << " and with exp convolutions." << endl;
  cout << "\t\tThe exp resolutions are read from config/expres.yaml";
  cout << "\n\n\t\twithout -m (default): calculate total C.S." << endl;
	cout << "\n\t\t-b: batch mode - calculate multiple configurations, ";
  cout << "e.g., in the case of inclusive measurements," << endl;
  cout << "\t\tIt can be done both for one nucleus (by, e.g., `-b rs.yaml`), or";
  cout << " for a batch of nuclei (by, e.g. `-b batch.yaml`)" << endl;
  cout << "\n\t\t-r: only calculate bound state Rl(r) and return, not to be used with -b";
  cout << "\n\t\t\033[1mNOTE: All the yaml files should be given without their paths." << endl;
  cout << "\t\tThe path should be written in config/basedir.yaml.\033[0m\n" << endl;
	exit(1);
} // end of member function PromptHelp

void TAMOMDIS::Go(){
  if(!fIsBatch){
    if(fIsMom) Parallel();
    else Total();
    return;
  } // end else

  // batch mode //
  const TAYaml &v = *PARA->GetPara();
  if(v["nus"]){ // super batch mode
    for(auto &t : v["nus"]){
      if(fIsMom) fnv.push_back(new TANuMom(t.as<string>(), fIsConv));
      else fnv.push_back(new TANucleus(t.as<string>()));
    } // end for
  } // end if
  else{
    // inferior batch mode //
    if(fIsMom) fnv.push_back(new TANuMom(PARA->cfg(), fIsConv));
    else fnv.push_back(new TANucleus(PARA->cfg()));
  } // end else
  // do the calculation //
  for(auto &t : fnv){
    if(fIsMom) dynamic_cast<TANuMom *>(t)->Mom();
    else t->Rs();
  } // end for
} // end member function Go

// -1: undefined; 1: true (zv: 1); 0: false (zv: 0)
bool TAMOMDIS::IsProtonRemoval() const{
  if(-1 == fIsProtonRemoval) TAException::Warn("TAMOMDIS", "IsProtonRemoval: not defined");
  if(0 != fIsProtonRemoval && 1 != fIsProtonRemoval)
    TAException::Warn("TAMOMDIS", "IsProtonRemoval: abnormal fIsProtonRemoval value: %d",
      fIsProtonRemoval);
  return fIsProtonRemoval;
} // end member function IsProtonRemoval

////------------ PREPARE FOR THE MOMDIS INTEGRAL ----------------/////
// solve the boundwave function for the valence nucleon and calculate the S-Matrix
void TAMOMDIS::Configure(){
  if(fHasConfigured) return;
  fHasConfigured = true;

  // to calculate the S-matrix and fit it for alphaj and betaj //
  const TAYaml &v = *PARA->GetPara();
  // f_NN parameterization: 0: horiuchi; 1: LenziRay
  if(v["FNNParOpt"]) TAFNN::SetParOption(v["FNNParOpt"].as<int>());
  const int zP = v["ZP"].as<int>(), aP = v["AP"].as<int>();
  const int zT = v["ZT"].as<int>(), aT = v["AT"].as<int>();
  const int zc = v["Zc"].as<int>(), ac = v["Ac"].as<int>();
  const int zv = zP - zc, av = aP - ac;
  if((zv != 0 && zv != 1) || av != 1) TAException::Error("TAMOMDIS",
    "Configure: abnormal zv(%d) or av(%d)", zv, av);
  fIsProtonRemoval = zv;
  const double Ek = v["Ek"].as<double>(); // in MeV/nucleon
  // register the reaction info //
  if("" == fBanner){
    std::stringstream ss; ss.clear(); ss.str("");
    ss << "\033[34;1m->->->- " << AME->name(zT, aT) << "(" << AME->name(zP, aP);
    ss << "," << AME->name(zc, ac) << ")" << AME->name(zv, av) << " @" << Ek;
    ss << " MeV/nucleon\033[0m -<-<-<-";
    fBanner = ss.str();
  } // end if

  // core's and target's nucleonic density distributions //
  // core density
  const string densC = v["densC"] ? v["densC"].as<string>() : "";
  // target density
  const string densT = v["densT"] ? v["densT"].as<string>() : "";
  // the optical optentials //
  // c+T optical potential
  const string optC = v["optC"] ? v["optC"].as<string>() : "";
  // v+T optical potential
  const string optV = v["optV"] ? v["optV"].as<string>() : "";
  // nucleon size factor: exp(-r^2/ap^2)
  const double alphav = v["alphav"].as<double>(), alphaC = v["alphaC"].as<double>();
  const double alphaT = v["alphaT"].as<double>();
  if(fExCS) return; // using external CS - skip bound and S-Matrices calculations
  fBound = new TABound(v); // to solve the radial wavefunction
  if(fIsBound) exit(1);
  if(densT != ""){ // t-rho-rho method
    fSc = new TASMatrix(zc, ac, zT, aT, Ek, densC, densT, alphaC, alphaT);
    fSv = new TASMatrix(zv, av, zT, aT, Ek,    "", densT, alphav, alphaT);
  } else{ // optical potential method
    fSc = new TASMatrix(zc, ac, zT, aT, Ek, optC, alphaC, alphaT);
    fSv = new TASMatrix(zv, av, zT, aT, Ek, optV, alphav, alphaT);
  } // end else
  l = fBound->Getl(); // the orbital angular momentum quantum number
  fMOM_M = new TAMOMDIS_M(this);

  // the convolver to generate the exp. mom distri. //
  if(fIsConv) fConv = new TAConvolve(aP, zP, zc, Ek, fExpResCfgFile);
} // end member function Configure

double TAMOMDIS::Total(){
  const int l1 = l+1;
  fSigmaStr_M.resize(l1); fSigmaDiff_M.resize(l1); fSigmaTotal_M.resize(l1);
  fSigmaStr_M.assign(l1,0.); fSigmaDiff_M.assign(l1,0.); fSigmaTotal_M.assign(l1,0.);
  fSigmaStr = 0.; fSigmaDiff = 0.; fSigmaTotal = 0.;
  // calculate total knockout X-sections //
  for(int m = 0; m <= l; m++) fSigmaStr_M[m] = fMOM_M->CSStr(m);
  for(int m = 0; m <= l; m++) fSigmaDiff_M[m] = fMOM_M->CSDiff(m);
  for(int m = 0; m <= l; m++){
    fSigmaTotal_M[m] = fSigmaStr_M[m] + fSigmaDiff_M[m];
    fSigmaStr += fSigmaStr_M[m];
    fSigmaDiff += fSigmaDiff_M[m];
    fSigmaTotal += fSigmaTotal_M[m];
  } // end for over m
  // print the results //
  cout << "\033[1mCalculating total knockout cross sections (in mb)\033[0m" << endl;
  cout << setprecision(6) << fixed; cout << endl;
  cout << "\033[34;1m" << setw(8) << "M" << setw(12);
  cout << "STRIP" << setw(12) << "DIFF" << setw(12) << " TOTAL" << "\033[0m" << endl;
  for(int m = 0; m <= l; m++){
    cout << setw(8) << m << setw(12) << fSigmaStr_M[m];
    cout << setw(12) << fSigmaDiff_M[m] << setw(12) << fSigmaTotal_M[m] << endl;
  } // end for over m
  cout << "In total:" << endl;
  cout << "Stripping c.s.:      " << setw(12) << fSigmaStr << " mb" << endl;
  cout << "Diffractive c.s.:    " << setw(12) << fSigmaDiff << " mb" << endl;
  cout << "Total knockout c.s.: " << setw(12) << fSigmaTotal << " mb" << endl;
  // cout << "OUTPUT TO FILE? <Ctrl-C> to cancel; <ENTER> to confirm..." << endl; getchar();
  // output to file //
  char tt[256]{}; strcpy(tt, TAException::time0(true)); // suitable for filename
  char tt1[256]; strcpy(tt1, TAException::time0()); // suitable for human reading
  char file[1024]; sprintf(file, "%s%s_%s.txt", PARA->BaseDir(), PARA->BaseName(), tt);
  ofstream fout(file); fout << std::setprecision(6) << std::fixed;
  fout << "Run at " << tt1 << endl;
  fout << "Orbit: " << fBound->Orbit() << endl;
  if((*PARA->GetPara())["Eref"])
    fout << "S_N+Ex: " << (*PARA->GetPara())["Eref"].as<double>() << " MeV" << endl;
  fout << endl << "The calculated total knockout cross sections (in mb)" << endl;
  fout << setw(8) << "M" << setw(12) << "STRIP";
  fout << setw(12) << "DIFF" << setw(12) << " TOTAL" << endl;
  for(int m = 0; m <= l; m++){
    fout << setw(8) << m << setw(12) << fSigmaStr_M[m];
    fout << setw(12) << fSigmaDiff_M[m] << " " << setw(12) << fSigmaTotal_M[m] << endl;
  } // end for over m
  fout << endl << "In total:" << endl;
  fout << "Stripping c.s.:      " << setw(12) << fSigmaStr << " mb" << endl;
  fout << "Diffractive c.s.:    " << setw(12) << fSigmaDiff << " mb" << endl;
  fout << "Total knockout c.s.: " << setw(12) << fSigmaTotal << " mb" << endl;
  cout << "The result is stored in \033[33;1m" << file << "\033[0m" << endl;
  fout.close();

  return fSigmaTotal;
} // end member function Total

///--- CALCULATE PARTIAL C.S. && CORE MOMENTUM DISTRIBUTION w.r.t. ---///
///--- valence nucleon in the projectile-incident direction ---///
void TAMOMDIS::Parallel(){
  cout << "\033[32;1mCalculating parallel momentum distributions.\033[0m" << endl;
  // calculate momdis and sigma for each m
  const int l1 = l+1;
  fSigmaStr_M.resize(l1); fSigmaDiff_M.resize(l1); fSigmaTotal_M.resize(l1);
  fMomStr.resize(fNmom); fMomDiff.resize(fNmom); fMomTotal.resize(fNmom);
  fMomStr.assign(fNmom,0.); fMomDiff.assign(fNmom,0.); fMomTotal.assign(fNmom,0.);
  // momentum distributions //
  fSigmaStr = 0.; fSigmaDiff = 0.; fSigmaTotal = 0.;
  fMomStr_M.resize(l1); fMomDiff_M.resize(l1); fMomTotal_M.resize(l1);
  for(int m = 0; m <= l; m++){ // resize 2D vectors
    fMomStr_M[m].resize(fNmom); fMomDiff_M[m].resize(fNmom); fMomTotal_M[m].resize(fNmom);
    fMomStr_M[m].assign(fNmom,0.); fMomDiff_M[m].assign(fNmom,0.); fMomTotal_M[m].assign(fNmom,0.);
  } // end for over i
  for(int m = 0; m <= l; m++) fSigmaStr_M[m] = fMOM_M->ParallelStr(m, fMomStr_M[m].data());
  for(int m = 0; m <= l; m++) fSigmaDiff_M[m] = fMOM_M->ParallelDiff(m, fMomDiff_M[m].data());
  for(int m = 0; m <= l; m++){
    // note that momdis integrated over kz is not as accurate as calcualted by Total() //
    fSigmaTotal_M[m] = fSigmaStr_M[m] + fSigmaDiff_M[m];
    fSigmaStr += fSigmaStr_M[m];
    fSigmaDiff += fSigmaDiff_M[m];
    fSigmaTotal += fSigmaTotal_M[m];
    for(int i = 0; i < fNmom; i++){
      fMomTotal_M[m][i] = fMomStr_M[m][i] + fMomDiff_M[m][i];
      fMomStr[i] += fMomStr_M[m][i];
      fMomDiff[i] += fMomDiff_M[m][i];
      fMomTotal[i] += fMomTotal_M[m][i];
    } // end for over i
  } // end for over m
  // output stripping momdis to file //
  char tt[256]{}; strcpy(tt, TAException::time0(true)); // suitable for filename
  char file0[1024];
  // sprintf(file, "%s%s_strT%s.txt", PARA->BaseDir(), PARA->BaseName(), tt);
  sprintf(file0, "%s%s_strT.txt", PARA->BaseDir(), PARA->BaseName());
  ofstream fout0(file0);
  for(int i = fNmom; i--;) fout0 << -kz(i) << " " << fMomStr[i] << endl;
  for(int i = 1; i < fNmom; i++) fout0 << kz(i) << " " << fMomStr[i] << endl;
  fout0.close();
  cout << "Total stripping distribution has been written to file " << file0 << endl;
  // output diffractive momdis to file //
  // sprintf(file, "%s%s_difT%s.txt", PARA->BaseDir(), PARA->BaseName(), tt);
  char file1[1024];
  sprintf(file1, "%s%s_difT.txt", PARA->BaseDir(), PARA->BaseName());
  ofstream fout1(file1);
  for(int i = fNmom; i--;) fout1 << -kz(i) << " " << fMomDiff[i] << endl;
  for(int i = 1; i < fNmom; i++) fout1 << kz(i) << " " << fMomDiff[i] << endl;
  fout1.close();
  cout << "Total diffractive distribution has been written to file " << file1 << endl;
  // output total momdis to file //
  // sprintf(file, "%s%s_tot%s.txt", PARA->BaseDir(), PARA->BaseName(), tt);
  char file2[1024];
  sprintf(file2, "%s%s_tot.txt", PARA->BaseDir(), PARA->BaseName());
  ofstream fout2(file2);
  for(int i = fNmom; i--;) fout2 << -kz(i) << " " << fMomTotal[i] << endl;
  for(int i = 1; i < fNmom; i++) fout2 << kz(i) << " " << fMomTotal[i] << endl;
  fout2.close();
  cout << "Total distribution has been written to file " << file2 << endl;


  // print the results //
  cout << setprecision(6) << fixed;
  cout << "\033[34;1m" << setw(8) << "M" << setw(12);
  cout << "STRIP" << setw(12) << "DIFF" << setw(12) << " TOTAL";
  cout << "\033[0m" << endl;
  for(int m = 0; m <= l; m++){
    cout << setw(8) << m << setw(12) << fSigmaStr_M[m];
    cout << setw(12) << fSigmaDiff_M[m] << setw(12) << fSigmaTotal_M[m] << endl;
  } // end for over m
  cout << "In total:" << endl;
  cout << "Stripping c.s.:      " << setw(12) << fSigmaStr << " mb" << endl;
  cout << "Diffractive c.s.:    " << setw(12) << fSigmaDiff << " mb" << endl;
  cout << "Total knockout c.s.: " << setw(12) << fSigmaTotal << " mb" << endl;

  // copy the momdis files to vis/ for quick visualization //
  const char *strT = "vis/strT.txt"; copy(file0, strT);
  const char *difT = "vis/difT.txt"; copy(file1, difT);
  const char *tot  = "vis/tot.txt";  copy(file2, tot);
  cout << "The generated momentum distributions can be visualized via ";
  cout << "`cd vis/` and `root -l mom.C`" << endl;

  if(!fIsConv) return;

  //// calculate the convoluted experimental momdis ////
  cout << "Convolving to give the experimental momentum distribution..." << endl;
  const int N = 2*fNmom-1;
  // XXX VLA: varialble length array; FIXME //
  double x[N], f[N]; for(int i = N; i--;) x[i] = f[i] = 0.;
  int j = 0;
  // construct the original str momdis arrays //
  for(int i = fNmom; i--;){ x[j] = -kz(i); f[j] = fMomStr[i]; j++; }
  for(int i = 1; i < fNmom; i++){ x[j] = kz(i); f[j] = fMomStr[i]; j++; }
  // do the convolution //
  char file3[1024];
  sprintf(file3, "%s%s_strTConv.txt", PARA->BaseDir(), PARA->BaseName());
  ofstream fout3(file3);
  Convolve(N, x, f, fSigmaStr, fMomConv, fMomStrConv, fout3);
  cout << "Str exp distribution has been written to file " << file3 << endl;
  // construct the original total momdis arrays //
  j = 0;
  for(int i = fNmom; i--;){ x[j] = -kz(i); f[j] = fMomTotal[i]; j++; }
  for(int i = 1; i < fNmom; i++){ x[j] = kz(i); f[j] = fMomTotal[i]; j++; }
  // do the convolution //
  // sprintf(file, "%s%s_totConv%s.txt", PARA->BaseDir(), PARA->BaseName(), tt);
  char file4[1024];
  sprintf(file4, "%s%s_totConv.txt", PARA->BaseDir(), PARA->BaseName());
  ofstream fout4(file4);
  Convolve(N, x, f, fSigmaTotal, fMomConv, fMomTotalConv, fout4);
  cout << "Total exp distribution has been written to file " << file4 << endl;
} // end the member function Parallel

// convolve the proj. c.m. momdis (m) to give the exp. distri. stored in fout
// cc: normalization constant; fconv: convolved f, return value of Convolve()
void TAMOMDIS::Convolve(int l, double *x, double *f, double cc, vector<double> &xconv,
  vector<double> &fconv, ofstream &fout){
  return fConv->ExpMomDis(l, x, f, cc, xconv, fconv, fout);
} // end member function conv

double TAMOMDIS::kz(int i) const{
  return fMOM_M->kz(i);
} // end member function kz

double TAMOMDIS::kzconv(int i) const{
  if(i < 0 || i >= int(fMomConv.size()))
    TAException::Error("TAMOMDIS", "kzconv: the required subscipt (i=%d) is invalid", i);
  return fMomConv[i];
} // end member function kzconv

int TAMOMDIS::nmomconv() const{
  if(!fMomConv.size()) TAException::Warn("TAMOMDIS", "nmomconv: 0 length of fMomConv");
  return int(fMomConv.size());
} // end member function nmomconv

void TAMOMDIS::testm(int &m) const{
  if(m < 0) m = -m;
  if(m > l) TAException::Error("TAMOMDIS", "testm: m=%d > l=%d", m, l);
} // end member function testm

TABound *TAMOMDIS::GetBound(){
  if(!fBound) Configure();
  return fBound;
} // end member function GetBound
double TAMOMDIS::Rl(double r){ return GetBound()->Rl(r); } ///< the radial wavefunction
double TAMOMDIS::u(double r){ return GetBound()->u(r); } ///< the radial wavefunction
/// returns the S-matrix of the core-target
TASMatrix *TAMOMDIS::Sc(){
  if(!fSc) Configure();
  return fSc;
} // end member function GetSMatrix
/// returns the S-matrix of the valence nucleon-target
TASMatrix *TAMOMDIS::Sv(){
  if(!fSv) Configure();
  return fSv;
} // end member function GetSMatrix
cdouble TAMOMDIS::Sc(double b){ return Sc()->SMatrix(b); }
cdouble TAMOMDIS::Sv(double b){ return Sv()->SMatrix(b); }

#ifdef CATCH_CNOK
// TEST TEST TEST TEST TEST TEST //
TEST_CASE("MOMD TEST", "[.momd]"){
  TAParaManager::Instance()->SetBaseDir("config/C/C16");
  TAMOMDIS *mom = new TAMOMDIS("B15.yaml");
  const double cs = mom->Total(); // calculate total one-nucleon knockout X-section
  CHECK(cs == Approx(17.475).epsilon(1e-2));
} // end TEST_CASE
#endif
